from langchain.agents import create_agent
from langchain_core.runnables import RunnableConfig
from langchain_deepseek import ChatDeepSeek
from langgraph.config import get_store
from langgraph.store.memory import InMemoryStore
from langchain_core.tools import tool

# 定义⻓期存储
store = InMemoryStore()
# 添加⼀些测试数据。 users是命名空间，user_123是key，后⾯的JSON数据是value
store.put(
    ("users",),
    "user_123",
    {
        "name": "张三",
        "age": "40",
    }
)


# 定义⼯具
@tool(return_direct=True)
def get_user_info(config: RunnableConfig) -> str:
    """查找⽤户信息"""
    # 获取⻓期存储。获取到了后，这个存储组件可读也可写
    store = get_store()

    # 获取配置中的⽤户ID
    user_id = config["configurable"].get("user_id")
    user_info = store.get(("users",), user_id)
    return str(user_info.value) if user_info else "Unknown user"


# 初始化DeepSeek模型客户端，配置模型参数
llm = ChatDeepSeek(
    model="deepseek-chat",  # 指定模型名称
    api_key="sk-dc3f063534c542a3a403760655d3eae2",  # 替换为你的API Key
    temperature=0.7,  # 控制生成结果的随机性，0表示更确定性的输出
    max_tokens=None,  # 生成内容的最大长度，None代表使用模型默认值
    timeout=60,  # 请求超时设置
    max_retries=2,  # 请求失败重试次数
)

agent = create_agent(
    model=llm,
    tools=[get_user_info],
    store=store  # 配置存储组件, 保存对话上下文
)
# Run the agent
res = agent.invoke(
    {"messages": [{"role": "user", "content": "查找⽤户信息"}]},
    config={"configurable": {"user_id": "user_123"}}
)

print(res)
